#%%
import qutip as qt
import numpy as np

paulis = {'X': qt.sigmax, 'Y': qt.sigmay, 'Z': qt.sigmaz}


def place(pois, n, base_gate='X'):
    # 在n比特的pois位置上作用basic_gate
    g = [qt.identity(2) for i in range(n)]
    for i in pois:
        g[i] = paulis[base_gate]()
    return qt.tensor(g)


def classical_solution(color, connect):
    res = 0
    for c in connect:
        if color[c[0]] != color[c[1]]:
            res += 1
    return res


def xy(p0, p1, n):
    return place([p0, p1], n, 'X') + place([p0, p1], n, 'Y')


def generate_hamiltonian(connect, n):
    H = 0
    for c in connect:
        H += xy(*c, n)
    return H


# %%
node = 10
all_connect = [(i, j) for i in range(node - 1) for j in range(i + 1, node)]
np.random.shuffle(all_connect)
num_connect = np.random.choice(range(1, node + 1))
connect = all_connect[:num_connect]
color = np.random.choice([0, 1], size=node)

# %%
# node = 2
# connect = [(0, 1)]
# color = [0, 1]
# %%
print('Node: ', color)
print('Connect: ', connect)
print("classical solution: ", classical_solution(color, connect))
# %%
init_state = qt.ket(color)
ham = generate_hamiltonian(connect, node)
res = qt.expect(ham*ham, init_state)/4
print(res)
# %%
